package com.flink.keyby_partition;

import com.flink.datasource.UserSource;
import com.flink.entity.User;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.util.Map;

/**
 * 描述:
 * 按键分区状态,模拟窗口，统计user访问次数
 *
 * @author yanzhengwu
 * @create 2022-08-30 23:43
 */
public class KeyByPartitionUserCount {

    public static void main(String[] args) throws Exception {
        //声明执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        //测试为了保证元素顺序并行度设置为1，可以理解为单线程执行
        env.setParallelism(1);


        SingleOutputStreamOperator<User> stream = env.addSource(new UserSource())
                //指定乱序流多等1秒
                .assignTimestampsAndWatermarks(WatermarkStrategy.<User>forBoundedOutOfOrderness(Duration.ofSeconds(1L))
                        .withTimestampAssigner(new SerializableTimestampAssigner<User>() {
                            @Override
                            public long extractTimestamp(User user, long recordTimestamp) {
                                return user.getTimestamp();
                            }
                        }));
        //stream.print("dataSource");
        stream.keyBy(User::getName)
                .process(new MyFakeWindowResult())
                .print("count");

        env.execute();
    }

    /**
     * TODO 先按键分区后再 存的状态就是每个键分区中的状态在进行 窗口分组 输出
     */
    public static class MyFakeWindowResult extends KeyedProcessFunction<String, User, String> {
        //定义窗口大小 10秒一次计算
        private Long windowSize = 10L;
        //定义map状态
        MapState<Long, Long> mapState;

        @Override
        public void processElement(User value, Context ctx, Collector<String> out) throws Exception {
            //整除计算窗口开始时间
            long windowStart = value.getTimestamp() / windowSize * windowSize;
            long windowEnd = windowStart + windowSize;
            //计算user访问次数
            if (mapState.contains(windowStart)) {
                long count = mapState.get(windowStart);
                mapState.put(windowStart, count + 1);
            } else {
                mapState.put(windowStart, 1L);
            }
            //定义10秒-1毫秒的定时器执行一次 模拟窗口计算
            ctx.timerService().registerEventTimeTimer(windowEnd - 1);


        }

        @Override
        public void onTimer(long timestamp, OnTimerContext ctx, Collector<String> out) throws Exception {
            //这个地方由于注册定时器 -1 所以需要 +1 补齐
            long windowEnd = timestamp + 1;
            long windowStart = windowEnd - windowSize;
            Long count = mapState.get(windowStart);
            out.collect(windowStart + "~" + windowEnd + " user:  " + ctx.getCurrentKey() + " count:" + count);

            //使用完窗口数据 及时回收避免内存爆炸
            mapState.remove(windowStart);
        }

        @Override
        public void open(Configuration parameters) throws Exception {
            mapState = getRuntimeContext().getMapState(new MapStateDescriptor<Long, Long>("map-state", Long.class, Long.class));
        }
    }
}
